# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
# ~~~~~~~ 使用LSB隐写术对图像中的秘密消息进行编码 ~~~~~~~ #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

from skimage.io import imread, imsave
import sys
from DataHidingDemo.extra import file_to_bin, generate_difference, mes_to_bin
from PIL import PngImagePlugin
import wave
import random
import numpy as np
import math
import random


# TODO: 更多的文件格式支持
# TODO: 随机位置隐写功能的支持  可以为不同算法：
#       混沌序列、
#       线性反馈移位寄存器随机数发生器、
#       python自带的伪随机数发生器
# TODO: 隐写提取效率的优化

# PSNR峰值信噪比
def psnr(img_1, img_2):
    mse = np.mean((img_1 / 1.0 - img_2 / 1.0) ** 2)
    if mse < 1.0e-10:
        return 100
    return 10 * math.log10(255.0 ** 2 / mse)

# 错误率测试
def falseRate(img1, img2):
    img1_float = np.array(img1, dtype=float)
    img2_float = np.array(img2, dtype=float)
    return np.mean(abs(img1_float-img2_float))/256

def lsb_hide(img, mes="", filename="", bits=1):
    # 保存原始图像
    print("保存原始图像")
    original = img.copy()

    # 找出图像的大小和频道数
    print("找出图像的大小和频道数")
    a, b = img.shape[0], img.shape[1]
    c = 3 if len(img.shape) == 3 else 1
    print("大小:", a, "X", b)
    print("频道数:", c)
    if mes == "":
        # 接收的秘密消息转化比特流
        print("接收的秘密消息转化比特流:")
        bitstream = file_to_bin(filename, bits)
        print(bitstream)
    else:
        # 接收的秘密消息转化比特流
        print("接收的秘密消息转化比特流:")
        
        bitstream = mes_to_bin(mes, bits)
        print(bitstream)

    # 检查图像的容量
    print("检查图像的容量")
    capacity = a * b * c;
    if capacity < len(bitstream):
        print("文件太大无法隐写!")
        return

    i, x = 0, 0
    while i < a and x < len(bitstream):
        for j in range(b):
            k = 0
            while k < c and x < len(bitstream):
                # 删除最后的位并替换为秘密消息位
                if c == 1:
                    img[i, j] += -(img[i, j] % (2 ** bits)) + int(bitstream[x], base=2)
                else:
                    img[i, j, k] += -(img[i, j, k] % (2 ** bits)) + int(bitstream[x], base=2)
                k += 1
                x += 1
        i += 1
    
    return img,original
    # 生成原始图像和输出之间的差异
    # if difference:
    #     generate_difference(original, img)
    #
    # 保存文件
    #imsave(out_name + '.png', img)
    
def wav_lsb_hide(wav, mes="", filename=""):
    cover_wav = wave.open(wav, mode='rb')
    frames = bytearray(cover_wav.readframes(cover_wav.getnframes()))

    message_bits = []
    mess_len = len(mes) * 8
    message_bits += [int(bit) for bit in format(mess_len, '016b')]
    # print(message_bits)
    for l in mes:
        message_bits += [(ord(l) >> i) & 1 for i in range(8)]
    # print(message_bits)
    for i in range(len(message_bits)):
        if frames[i] % 2 != message_bits[i]:
            if frames[i] == 255:
                s = -1
            elif frames[i] == 0:
                s = +1
            else:
                s = random.choice([-1, +1])
            frames[i] = frames[i] + s

    cover_wav.close()
    return frames


def main():
    print("Enter the source path of the image to encode:")

    file = input()
    img = imread(file)
    print("Enter message to be hidden: ")
    file2 = input()
    lsb_hide(img, file2, 1, 'output', False)


if __name__ == '__main__':
    main()
